package cn.zengdeai.langchain4j_springboot.config;


import cn.zengdeai.langchain4j_springboot.service.ToolsService;
import dev.langchain4j.community.model.dashscope.QwenEmbeddingModel;
import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.memory.chat.ChatMemoryProvider;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.chat.StreamingChatLanguageModel;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.service.*;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

@Configuration
public class AiConfig {

    public int counter=0;

    @Autowired
    public QwenEmbeddingModel qwenEmbeddingModel;


    public interface Assistant {
        String chat(String message);
        // 流式响应
        TokenStream stream(String message);

        @SystemMessage("""
                您是“Tuling”航空公司的客户聊天支持代理。请以友好、乐于助人且愉快的方式来回复。
                        您正在通过在线聊天系统与客户互动。
                        在提供有关预订或取消预订的信息之前，您必须始终从用户处获取以下信息：预订号、客户姓名。
                        请讲中文。
			   今天的日期是 {{current_date}}.
        """)
        TokenStream stream(@UserMessage String message,@V("current_date") String currentDate);
    }



    public interface AssistantForCSR {
        @SystemMessage("""
                你现在模拟一个购买衣服的客户对我进行提问，我是一名电商客服，以训练并提高我的技能。问完一个问题问后再问下一个，不要一下子扔10几问题，不要总结。
			   今天的日期是 {{current_date}}.
        """)
        TokenStream streamForTrainForCSR(@UserMessage String message,@V("current_date") String currentDate);
    }

    public interface AssistantUnique {
        String chat(@MemoryId int memoryId, @UserMessage String userMessage);
        // 流式响应
        TokenStream stream(@MemoryId int memoryId, @UserMessage String userMessage);
    }


    @Bean
    public EmbeddingStore embeddingStore(){
        return new InMemoryEmbeddingStore();
    }


    /**
     * 带toolservice
     * @param qwenChatModel
     * @param qwenStreamingChatModel
     * @param embeddingStore
     * @param toolsService
     * @return
     */
    @Bean
    public Assistant assistant(ChatLanguageModel qwenChatModel,
                               StreamingChatLanguageModel qwenStreamingChatModel,
                               EmbeddingStore embeddingStore,
                               ToolsService toolsService) {
        ChatMemory chatMemory = MessageWindowChatMemory.withMaxMessages(10);


        ContentRetriever contentRetriever = EmbeddingStoreContentRetriever.builder()
                .embeddingStore(embeddingStore)
                .embeddingModel(qwenEmbeddingModel)
                .maxResults(5) // 最相似的5个结果
                .minScore(0.6) // 只找相似度在0.6以上的内容
                .build();

        Assistant assistant = AiServices.builder(Assistant.class)
                .chatLanguageModel(qwenChatModel)
                .streamingChatLanguageModel(qwenStreamingChatModel)
                .contentRetriever(contentRetriever)
                .tools(toolsService)
                .chatMemory(chatMemory)
                .build();
        System.out.println("计数器为:"+this.counter++);

        return  assistant;
    }


    /**
     * 不带toolservice
     * @param qwenChatModel
     * @param qwenStreamingChatModel
     * @return
     */
    @Bean
    public AssistantForCSR assistantCSR(ChatLanguageModel qwenChatModel,
                                     StreamingChatLanguageModel qwenStreamingChatModel) {
        ChatMemory chatMemory = MessageWindowChatMemory.withMaxMessages(10);
        AssistantForCSR assistantForCSR = AiServices.builder(AssistantForCSR.class)
                .chatLanguageModel(qwenChatModel)
                .streamingChatLanguageModel(qwenStreamingChatModel)
                .chatMemory(chatMemory)
                .build();
        return  assistantForCSR;
    }






    @Bean
    public AssistantUnique assistantUnique(ChatLanguageModel qwenChatModel,
                                           StreamingChatLanguageModel qwenStreamingChatModel) {

        AssistantUnique assistant = AiServices.builder(AssistantUnique.class)
                .chatLanguageModel(qwenChatModel)
                .streamingChatLanguageModel(qwenStreamingChatModel)
                .chatMemoryProvider(memoryId ->
                        MessageWindowChatMemory.builder().maxMessages(10)
                                .id(memoryId).build()
                )
                .build();

        return assistant;
    }

    @Bean
    public AssistantUnique assistantUniqueStore(ChatLanguageModel qwenChatModel,
                                                StreamingChatLanguageModel qwenStreamingChatModel) {

        PersistentChatMemoryStore store = new PersistentChatMemoryStore();

        ChatMemoryProvider chatMemoryProvider = memoryId -> MessageWindowChatMemory.builder()
                .id(memoryId)
                .maxMessages(10)
                .chatMemoryStore(store)
                .build();

        AssistantUnique assistant = AiServices.builder(AssistantUnique.class)
                .chatLanguageModel(qwenChatModel)
                .streamingChatLanguageModel(qwenStreamingChatModel)
                .chatMemoryProvider(chatMemoryProvider)
                .build();
        return assistant;
    }

}